import os
import shutil
import numpy as np

# 计算 IOU 的函数
def calculate_iou(box1, box2):
    # box = [x_center, y_center, width, height]
    x1_min = box1[0] - box1[2] / 2
    x1_max = box1[0] + box1[2] / 2
    y1_min = box1[1] - box1[3] / 2
    y1_max = box1[1] + box1[3] / 2

    x2_min = box2[0] - box2[2] / 2
    x2_max = box2[0] + box2[2] / 2
    y2_min = box2[1] - box2[3] / 2
    y2_max = box2[1] + box2[3] / 2

    # 计算交集区域
    inter_x_min = max(x1_min, x2_min)
    inter_y_min = max(y1_min, y2_min)
    inter_x_max = min(x1_max, x2_max)
    inter_y_max = min(y1_max, y2_max)

    if inter_x_max <= inter_x_min or inter_y_max <= inter_y_min:
        return 0.0

    inter_area = (inter_x_max - inter_x_min) * (inter_y_max - inter_y_min)

    # 计算并集区域
    box1_area = (x1_max - x1_min) * (y1_max - y1_min)
    box2_area = (x2_max - x2_min) * (y2_max - y2_min)
    union_area = box1_area + box2_area - inter_area

    return inter_area / union_area

# 主函数
def filter_txt_by_iou(input_dir, output_dir):
    
    # 阈值大就是重合的
    iou_thread = 0.95
    
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    if not input_dir.endswith('.txt'):
        print("input file is not txt !!!")
        return

    with open(input_dir, 'r') as f:
        lines = f.readlines()

    boxes_person = []
    boxes_helmet = []

    # 解析文件中的框信息
    for line in lines:
        parts = line.strip().split()
        class_id = int(parts[0])
        x, y, w, h = map(float, parts[1:])

        if class_id == 2:  # Person
            boxes_person.append([x, y, w, h])
        elif class_id == 3:  # Helmet
            boxes_helmet.append([x, y, w, h])
            
    # 检查 IOU 是否大于 0.9
    found_high_iou = False
    for person_box in boxes_person:
        for helmet_box in boxes_helmet:
            iou = calculate_iou(person_box, helmet_box)
            # print(iou)
            if iou > iou_thread:
                found_high_iou = True
                break
        if found_high_iou:
            break

    # 如果找到符合条件的文件，复制到输出目录
    if found_high_iou:
        save_name = os.path.join(output_dir, os.path.basename(input_dir))
        shutil.copy(input_dir, save_name)
    return found_high_iou
        
            
            
def read_file_paths(txt_file_path):
    """
    从指定的 txt 文件中读取文件路径。
    
    Args:
        txt_file_path (str): txt 文件的路径。
        
    Returns:
        list: 包含所有文件路径的列表。
    """
    file_paths = []
    
    try:
        with open(txt_file_path, 'r') as file:
            file_paths = [line.strip() for line in file if line.strip()]  # 去掉空行和两端的空格
    except FileNotFoundError:
        print(f"文件 {txt_file_path} 未找到！")
    except Exception as e:
        print(f"读取文件时发生错误: {e}")
    
    return file_paths



# 实现一个计算gt中是否出现“安全帽”和“人”标签几乎重提的情况；
# 如果存在上面情况就将txt标签筛选出来,并将txt的绝对路径保存到log.txt中
# iou_thread = 0.95 可以自行设置

# 从train.txt中获取训练数据（jpg）路径
txt_file_path = "/root/yjh/working_code/select_head_data/train-20241124-139781.txt"
file_paths = read_file_paths(txt_file_path)
# 数据保存文件夹路径
save_txt_label_floder = "/root/yjh/working_code/select_head_data/select_results"
# 遍历数据
current_path = os.getcwd()
log_file_path = os.path.join(current_path, 'log.txt')
with open(log_file_path, 'w') as log_file:
    for path in file_paths:
        base, _ = os.path.splitext(path)
        label_path = f"{base}.txt"
        # print(label_path)
        if os.path.exists(label_path): 
            save_switch = filter_txt_by_iou(label_path,save_txt_label_floder)
            if save_switch:
                log_file.write(path + '\n')
        else:
            print(f"路径{label_path}不存在！！！")
    
    
